import numpy as np
import openml


classification_tasks = [
    232, 236, 241, 245, 253, 254, 256, 258, 260, 262, 267, 271, 273, 275, 279, 288, 336,
    340, 2119, 2120, 2121, 2122, 2123, 2125, 2356, 3044, 3047, 3048, 3049, 3053, 3054,
    3055, 75089, 75092, 75093, 75098, 75100, 75108, 75109, 75112, 75114, 75115, 75116,
    75118, 75120, 75121, 75125, 75126, 75129, 75131, 75133, 75134, 75136, 75139, 75141,
    75142, 75143, 75146, 75147, 75148, 75149, 75153, 75154, 75156, 75157, 75159, 75161,
    75163, 75166, 75169, 75171, 75173, 75174, 75176, 75178, 75179, 75180, 75184, 75185,
    75187, 75192, 75195, 75196, 75199, 75210, 75212, 75213, 75215, 75217, 75219, 75221,
    75223, 75225, 75232, 75233, 75234, 75235, 75236, 75237, 75239, 75250, 126021, 126024,
    126028, 126030, 126031, 146574, 146575, 146576, 146577, 146578, 146583, 146586,
    146592, 146593, 146594, 146596, 146597, 146600, 146601, 146602, 146603, 146679,
    166859, 166866, 166872, 166875, 166882, 166897, 166905, 166906, 166913, 166915,
    166931, 166932, 166944, 166950, 166951, 166953, 166956, 166957, 166958, 166959,
    166970, 166996, 167085, 167086, 167087, 167088, 167089, 167090, 167094, 167096,
    167097, 167099, 167100, 167101, 167103, 167105, 167106, 167202, 167203, 167204,
    167205, 168785, 168791, 189779, 189786, 189828, 189829, 189836, 189840, 189841,
    189843, 189844, 189845, 189846, 189857, 189858, 189859, 189863, 189864, 189869,
    189870, 189875, 189878, 189880, 189881, 189882, 189883, 189884, 189887, 189890,
    189893, 189894, 189899, 189900, 189902, 190154, 190155, 190156, 190157, 190158,
    190159, 211720, 211721, 211722, 211723, 211724
]
regression_tasks = [
    359997, 359998, 359999, 360000, 360001, 360002, 360003, 167146, 360004, 360005, 360006,
    360007, 211696, 360009, 360010, 360011, 360012, 360013, 360014, 360015, 360016, 360017,
    360018, 360019, 360020, 360021, 360022, 360023, 360024, 360025, 360026, 360027, 360028,
    360029, 360030, 360031, 360032, 360033, 360034, 360035, 360036, 360037, 360038, 360039,
    360040, 360041, 360042, 360043, 360044, 360045, 360046, 360047, 360048, 360049, 360050,
    360051, 360052, 360053, 360054, 360055, 360056, 360057, 360058, 360059, 360060, 360061,
    360062, 360063, 360064, 360066, 360067, 360068, 360069, 360070, 360071, 360072, 360073,
    360074, 360075, 360076, 360077, 360078, 360079, 360080, 360081, 360082, 360083, 360084,
    360085, 360086, 360087, 360088, 360089, 360090, 360091, 360092, 360093, 360094, 360095,
    360096, 360097, 360098, 360100, 360101, 360102, 360103, 360104, 360105, 360106, 360107,
    360108,
]


def load_task(task_id):
    task = openml.tasks.get_task(task_id)
    X, y = task.get_X_and_y()
    train_indices, test_indices = task.get_train_test_split_indices()
    X_train = X[train_indices]
    y_train = y[train_indices]
    X_test = X[test_indices]
    y_test = y[test_indices]
    dataset = openml.datasets.get_dataset(task.dataset_id)
    _, _, cat, _ = dataset.get_data(target=task.target_name)
    name = dataset.name.lower()
    del _
    del dataset
    cat = ['categorical' if c else 'numerical' for c in cat]

    if isinstance(task, openml.tasks.OpenMLClassificationTask):
        task_type = 'classification'
    elif isinstance(task, openml.tasks.OpenMLRegressionTask):
        task_type = 'regression'
    else:
        raise ValueError('Unknown task type')

    return X_train, y_train, X_test, y_test, cat, task_type, name
